Henry Lee (이현규)

You Only Look Once: Unified, Real-Time Object Detection

2015

Joseph Redmon∗, Santosh Divvala∗†, Ross Girshick¶, Ali Farhadi∗†
University of Washington∗, Allen Institute for AI†, Facebook AI Research¶

YOLOv1 논문의 핵심은
  1. 이미지를 \( S \times S \) 그리드로 나누고
  2. 각 셀은 \( 5 \times B + C \) 인 예측한다. \(5\) 는 Confidence score, centerX(cX), centerY(cY), width(w), height(h)를 의미하고, \(B\) 는 박스 개수 (논문에서 2개), \(C\) 는 분류 클래스 개수 (논문은 PASCAL VOC 데이터 기반으로 20개)를 의미한다.
  3. 따라서 YOLO는 98개(\(7\times7\times B\))개의 박스를 예측한다. 하지만 클래스 분류는 각 박스에 할당하지 않고 셀 별(per cell) 예측한다.
  4. 결과적으로 Ground Truth에서 Object의 cX, cY가 위치하는 셀에 대하여 Prediction의 동일 셀은 2개의 박스 정보와 1개의 클래스 정보를 가지는데, 2개 박스 중 IoU가 높은 박스 값을 이용하여 위치 예측(Localization)을 학습한다.
  5. Confidence score는 Localization에서 사용되지 않은 박스에는 가중치(\(\lambda_{noobj}=0.5\))를 주어서 셀 내에 Object가 있는 지 없는 지 구분할 수 있도록 학습한다. Localization에 사용된 박스는 \(IoU^{GT}_{Pred}\)에 대하여 학습하여 박스 예측의 정확도 또한 Confidence에 반영하도록 한다.
  6. Localization과 동일하게, Ground Truth에서 Object의 cX, cY가 위치하는 셀에 대하여, 클래스 분류 예측(Classification)을 학습한다.

본격적인 구현에 앞서, 라이브러리를 불러온다.

  • numpy
  • tensorflow

GPU 리소스를 처음부터 full로 먹지 않도록 set_memory_growth() 를 사용한다. 

                                
import numpy as np
import tensorflow as tf

print(tf.__version__)
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
    tf.config.experimental.set_memory_growth(gpus[0], True)
except RuntimeError as e:
    # 프로그램 시작시에 메모리 증가가 설정되어야만 합니다
    print(e)
논문과 달리, 텐서플로에서 제공하는 "ImageNet 데이터셋으로 사전 훈련된(Pre-trained) VGG16"을 백본(Backbone)으로 불러온다.
이때, 입력 이미지 사이즈는 논문과 동일하게 448x448로 한다.
활성함수(activation function)로 마지막 레이어에서 Linear를 사용하고, 나머지 모든 레이어는 LeakyRelu를 사용한다. 
                                
backbone = tf.keras.applications.vgg16.VGG16(
    include_top=False,
    weights='imagenet',
    input_shape=(448,448,3)
)

for layer in backbone.layers:
    layer.trainable = False
    try:
        layer.activation = tf.keras.layers.LeakyReLU(alpha=0.1)
    except:
        pass
드롭아웃(Dropout) 레이어는 비율을 0.5로 하여 첫 FC 레이어 뒤에 둔다.
학습(training)에서 Loss가 튀는 것을 방지하기 위해서 Linear 대신 Sigmoid를 사용한다.
최종 출력 shape은 7x7x30이다. 
                                
x = tf.keras.layers.Conv2D(
    filters=1024,
    kernel_size=3,
    padding='same',
    activation=tf.keras.layers.LeakyReLU(alpha=0.1),
)(backbone.output)
x = tf.keras.layers.Conv2D(
    filters=1024,
    kernel_size=3,
    strides=(2,2),
    padding='same',
    activation=tf.keras.layers.LeakyReLU(alpha=0.1),
)(x)

x = tf.keras.layers.Conv2D(
    filters=1024,
    kernel_size=3,
    padding='same',
    activation=tf.keras.layers.LeakyReLU(alpha=0.1),
)(x)
x = tf.keras.layers.Conv2D(
    filters=1024,
    kernel_size=3,
    padding='same',
    activation=tf.keras.layers.LeakyReLU(alpha=0.1),
)(x)

x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(
    4096,
    activation=tf.keras.layers.LeakyReLU(alpha=0.1),
)(x)
x = tf.keras.layers.Dropout(0.5)(x)

x = tf.keras.layers.Dense(
    1470,
    activation='sigmoid',
)(x)
x = tf.keras.layers.Reshape((7,7,30))(x)

yolo = tf.keras.Model(inputs=backbone.input, outputs=x, name='Yolo')

손실함수(loss function)를 계산하기 위해서, 먼저 IoU 값을 구하는 function을 구현해야 한다.

손실 함수는

Classification Loss

\( \displaystyle\sum_{i=0}^{S^2} \Bbb{1}^{obj}_{i} \sum_{c\in class} (p_i(c)-\hat{p_i}(c))^2 \)

Confidence Loss

\( \displaystyle\sum_{i=0}^{S^2} \sum_{j=0}^{B} \Bbb{1}^{obj}_{ij} (C_i-\hat{C_i})^2 + \lambda_{noobj}\displaystyle\sum_{i=0}^{S^2} \sum_{j=0}^{B} \Bbb{1}^{noobj}_{ij} (C_i-\hat{C_i})^2 \)

Box Loss

\( \lambda_{coord} \displaystyle\sum_{i=0}^{S^2} \sum_{j=0}^{B} \Bbb{1}^{obj}_{ij} [ (x_i-\hat{x}_i)^2 + (y_i-\hat{y}_i)^2 + (\sqrt{w_i}-\sqrt{\hat{w}_i})^2 + (\sqrt{h_i}-\sqrt{\hat{h_i}})^2 ] \)

여기서 \( \Bbb{1}^{obj}_{i} \)의 의미는 Ground Truth(7x7x25)에서 confidence score가 1.인 경우,
\( \Bbb{1}^{obj}_{ij} \)의 의미는 Ground Truth(7x7x25)에서 confidence score가 1.이고, Prediction (7x7x30)에서 두 박스(B) 중 IoU가 더 높은 경우를 말한다.

즉, Classification Loss는 Ground Truth의 Confidence score가 1.인 경우 20개 클래스 예측값 차이 제곱의 총합.

Confidence Loss는 Ground Truth의 Confidence score가 1.이고 두 박스 중 IoU가 더 큰 박스의 Confidence score 차이 제곱과, 그렇지 않은 경우의 Confidence score 차이 제곱 x \( \lambda_{noobj}=0.5 \)의 총합.

Box Loss는 Ground Truth의 Confidence score가 1.이고 두 박스 중 IoU가 더 큰 박스의 x좌표 차이 제곱, y좌표 차이 제곱, 스퀘어루트한 w값 차이 제곱, 스퀘어루트한 h값 차이 제곱의 총합.
너비(w), 높이(h)에 스퀘어루트를 하는 이유는 큰 박스의 작은 오차보다 작은 상자의 오차에 가중치를 더 두기 위해서다. 

                                
h_grid, w_grid = 7, 7

def get_iou(box1, box2):
    h_offset, w_offset = np.indices((h_grid, w_grid))

    x1, y1, w1, h1 = tf.unstack(box1, num=4, axis=-1) 
    area_box1 = w1 * h1

    x1_global = (w_offset + x1) / w_grid
    y1_global = (h_offset + y1) / h_grid
    x1_min, y1_min = x1_global - w1/2, y1_global - h1/2
    x1_max, y1_max = x1_global + w1/2, y1_global + h1/2

    x2, y2, w2, h2 = tf.unstack(box2, num=4, axis=-1)
    area_box2 = w2 * h2

    x2_global = (w_offset + x2) / w_grid
    y2_global = (h_offset + y2) / h_grid
    x2_min, y2_min = x2_global - w2/2, y2_global - h2/2
    x2_max, y2_max = x2_global + w2/2, y2_global + h2/2

    x_i_min = tf.math.maximum(x1_min, x2_min)
    y_i_min = tf.math.maximum(y1_min, y2_min)
    x_i_max = tf.math.minimum(x1_max, x2_max)
    y_i_max = tf.math.minimum(y1_max, y2_max)

    w_i = x_i_max - x_i_min
    h_i = y_i_max - y_i_min
    w_i = tf.math.maximum(0., w_i)
    h_i = tf.math.maximum(0., h_i)
    area_i = w_i * h_i

    union = area_box1 + area_box2 - area_i + 1e-8
    return tf.expand_dims((area_i / union), axis=-1)

def loss_function(true, pred, coord=5.0, noobj=0.5):
    true_conf = true[..., :1]
    true_boxes = true[..., 1:5]
    true_classes = true[..., 5:]
    
    box1_conf = pred[..., :1]
    box1_boxes = pred[..., 1:5]
    box2_conf = pred[..., 5:6]
    box2_boxes = pred[..., 6:10]
    pred_classes = pred[..., 10:]

    # class, shape = (batch, 7, 7, 20)
    cls_loss = tf.math.square(pred_classes-true_classes)
    cls_loss = tf.where(
        tf.math.equal(true_conf, 1.0),
        cls_loss,
        0.0
    )
    cls_loss = tf.math.reduce_sum(cls_loss, axis=[1,2,3])

    # boxes, shape = (batch, 7, 7, 4)
    xy_loss = tf.math.square(
        box1_boxes[..., 0:2]-true_boxes[..., 0:2]
    )
    wh_loss = tf.math.square(
        tf.math.sqrt(box1_boxes[..., 2:4])-tf.math.sqrt(true_boxes[..., 2:4])
    ) # 작은 상자의 오차에 가중치를 더 두기 위해서 제곱근(square root)을 씌워 예측
    box1_loss = xy_loss + wh_loss
    box1_loss = tf.where(
        tf.math.equal(true_conf, 1.0),
        box1_loss,
        0.0
    )

    xy_loss = tf.math.square(
        box2_boxes[..., 0:2]-true_boxes[..., 0:2]
    )
    wh_loss = tf.math.square(
        tf.math.sqrt(box2_boxes[..., 2:4])-tf.math.sqrt(true_boxes[..., 2:4])
    )
    box2_loss = xy_loss + wh_loss
    box2_loss = tf.where(
        tf.math.equal(true_conf, 1.0),
        box2_loss,
        0.0
    )

    iou_box1 = get_iou(true_boxes, box1_boxes)
    iou_box2 = get_iou(true_boxes, box2_boxes)

    box1_loss = tf.where(
        tf.math.greater_equal(iou_box1, iou_box2), 
        box1_loss*coord, 
        0.
    )
    box2_loss = tf.where(
        tf.math.less(iou_box1, iou_box2), 
        box2_loss*coord, 
        0.
    )
    box_loss = box1_loss + box2_loss
    box_loss = tf.math.reduce_sum(box_loss, axis=[1,2,3])

    # confidence, shape = (batch, 7, 7, 1)
    box1_conf_loss = tf.math.square(box1_conf - true_conf*iou_box1)
    box1_conf_loss = tf.where(
        tf.math.greater_equal(iou_box1, iou_box2), 
        box1_conf_loss, 
        box1_conf_loss*noobj
    )

    box2_conf_loss = tf.math.square(box2_conf - true_conf*iou_box2)
    box2_conf_loss = tf.where(
        tf.math.less(iou_box1, iou_box2), 
        box2_conf_loss, 
        box2_conf_loss*noobj
    )

    conf_loss = box1_conf_loss + box2_conf_loss
    conf_loss = tf.math.reduce_sum(conf_loss, axis=[1,2,3])

    return tf.math.reduce_mean(cls_loss+box_loss+conf_loss)
학습하는 동안 배치 사이즈는 64, Optimizer의 모멘텀은 0.9, Decay는 0.0005로 한다. 
                                
sgd = tf.keras.optimizers.legacy.SGD(
    learning_rate=1e-3,
    momentum=0.9,
    decay=5e-4
)

yolo.compile(
    sgd,
    loss_function,
)
학습률(Learning Rate) 스케쥴은
75 epochs까지 0.001에서 0.01로 점차 증가 시키고,
105 epochs까지 0.001,
135 epochs까지 0.0001로 한다. 
                                
def scheduler(epoch, lr):
    if epoch < 75:
        return 1e-3 + 9e-3 * (epoch/75)
    elif epoch < 105:
        return 1e-3
    else:
        return 1e-4

callback = tf.keras.callbacks.LearningRateScheduler(scheduler)

yolo.fit(
    train_ds,
    epochs=135,
    validation_data=test_ds,
    callbacks=[callback],
)

Test Dataset에 대한 1 Batch 추론(Inference) 결과

Ground Truth: 파랑 (분류명과 중심점 표시)
Prediction: 빨강 (분류명과 Confidence score 그리고 박스 표시)